package com.hhy.lombok.processor;

import com.hhy.lombok.annotation.Proxy;
import com.hhy.lombok.proxy.OriginCaller;
import com.hhy.lombok.utils.JCUtils;
import com.hhy.lombok.utils.ProcessUtils;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.code.Symbol;
import com.sun.tools.javac.code.Type;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.ListBuffer;

import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import java.lang.annotation.Annotation;
import java.util.List;
import java.util.Set;

import static com.hhy.lombok.constant.Constant.*;
import static com.sun.tools.javac.util.List.nil;
import static com.sun.tools.javac.util.List.of;

/**
 * Created with IntelliJ IDEA.
 *
 * @author hai-yang
 * @date 2022/04/20/23:04
 * @description
 */

@SupportedAnnotationTypes("com.hhy.lombok.annotation.Proxy")
@SupportedSourceVersion(SourceVersion.RELEASE_8)
public class ProxyProcessor extends BaseProcessor {


    @Override
    public void doProcess(Element element) {
        JCTree jcTree = trees.getTree(element);
        jcTree.accept(new TreeTranslator() {
            @Override
            public void visitMethodDef(JCTree.JCMethodDecl jcMethodDecl) {
                Symbol.MethodSymbol methodSymbol = (Symbol.MethodSymbol) element;
                JCTree tree = getTree(methodSymbol.owner.type.toString());
                treeMaker.at(tree.pos);
                // 拿到注解
                List<Object> targets = (List) ProcessUtils.getAnnotationValue(element, Proxy.class, PROXY_ANNOTATION_TARGET_ATTR_NAME);
                int index = targets.size();
                for (Object target : targets) {
                    // 被代理的类型
                    Type.ClassType classType = (Type.ClassType) target;
                    enhanceMethod(jcMethodDecl, classType, --index);
                }
            }
        });
    }

    private void enhanceMethod(JCTree.JCMethodDecl jcMethodDecl, Type.ClassType classType, int index) {
        boolean last = index == 0;
        log.log("current index: ", index + "");
        ListBuffer finalMethod = new ListBuffer();
        // 原方法返回类型
        JCTree.JCExpression restype = jcMethodDecl.restype;
        boolean haveReturn = !"void".equals(restype.type.toString());

        JCTree.JCBlock newBody = jcMethodDecl.body;
        if (!haveReturn) {
            newBody = returnNullBody(newBody, 0);
        }

        // lambda 转为接口
        JCTree.JCLambda lambda = treeMaker.Lambda(nil(), newBody);
        JCTree.JCVariableDecl caller = createCaller(ProcessUtils.generateCallerVarName(index), lambda);
        finalMethod.add(caller);

        // 生成 代理对象
        // ProxyImpl p = new ProxyImpl();
        JCTree.JCVariableDecl proxyObj = JCUtils.createVar(treeMaker,
                names.fromString(ProcessUtils.generateProxyObjVarName(index)),
                JCUtils.memberAccess(treeMaker, elementUtils, classType.toString()),
                createProxyInstance(classType), Flags.PARAMETER);
        finalMethod.add(proxyObj);

        // 调用实现类的 process函数
        JCTree.JCExpressionStatement processor = generateProxyCallMethod(ProcessUtils.generateProxyObjVarName(index),
                of(treeMaker.Ident(elementUtils.getName(ProcessUtils.generateCallerVarName(index)))));

        // 生成 Object result = processor.ret
        JCTree.JCVariableDecl result = createResultVar(processor.expr);
        ListBuffer tryBlockList = finalMethod;
        if (last) {
            tryBlockList = new ListBuffer();
        }
        tryBlockList.add(result);

        if (haveReturn) {
            JCTree.JCReturn jcReturn = JCUtils.createReturn(treeMaker, treeMaker.Ident(result.name), null, restype);
            tryBlockList.add(jcReturn);
        }

        // 只有最外层的执行 才需要 try catch 去转换异常类型
        if (last) {
            JCTree.JCTry jcTry = tryCatchAndCastException(tryBlockList.toList(), jcMethodDecl.getThrows());
            finalMethod.add(jcTry);
        }

        jcMethodDecl.body = treeMaker.Block(0, finalMethod.toList());
        log.log("current index: ", index + ", method name: (", jcMethodDecl.name.toString(), ") [", jcMethodDecl.params.toString(), "] ",
                "enhanceMethod after: " + "\n" + jcMethodDecl.body);
    }

    private JCTree.JCTry tryCatchAndCastException(com.sun.tools.javac.util.List tryBlockList, com.sun.tools.javac.util.List<JCTree.JCExpression> jcThrows) {
        ListBuffer<JCTree.JCStatement> catchBody = new ListBuffer();
        for (JCTree.JCExpression jcThrow : jcThrows) {
            // 如果捕获的异常是需要抛出的： if(e instanceof Exception){ throw (Exception) e }
            JCTree.JCStatement ifTypeAndThrow = treeMaker.If(
                    treeMaker.TypeTest(treeMaker.Ident(elementUtils.getName(PROXY_EXCEPTION_VAR_NAME)), jcThrow),
                    treeMaker.Throw(treeMaker.TypeCast(jcThrow, treeMaker.Ident(elementUtils.getName(PROXY_EXCEPTION_VAR_NAME)))), // 强转成一样的异常
                    null);
            catchBody.add(ifTypeAndThrow);
        }

        //强转成 RuntimeException 类型 兜底策略
        JCTree.JCTypeCast jcTypeCastRuntimeException = treeMaker.TypeCast(memberAccess("java.lang.RuntimeException"),
                treeMaker.Ident(elementUtils.getName(PROXY_EXCEPTION_VAR_NAME)));
        //抛出 RuntimeException 类型的异常
        catchBody.add(treeMaker.Throw(jcTypeCastRuntimeException));

        // catch(java.lang.Throwable e) { }
        JCTree.JCVariableDecl exceptionVar = createVar(PROXY_EXCEPTION_VAR_NAME, "java.lang.Throwable", null);
        return treeMaker.Try(treeMaker.Block(0, tryBlockList), of(treeMaker.Catch(exceptionVar, treeMaker.Block(0, catchBody.toList()))), null);
    }

    private JCTree.JCVariableDecl createResultVar(JCTree.JCExpression statement) {
        return JCUtils.createVar(treeMaker,
                names.fromString(PROXY_RESULT_VAR_NAME), memberAccess(PROXY_RESULT_TYPE), statement, Flags.PARAMETER);
    }

    private JCTree.JCExpressionStatement generateProxyCallMethod(String instanceName, com.sun.tools.javac.util.List<JCTree.JCExpression> params) {
        return JCUtils.callMethod(treeMaker, elementUtils, instanceName, PROXY_PROCESS_NAME, params);
    }

    private JCTree.JCExpression createProxyInstance(Type.ClassType classType) {
        return treeMaker.NewClass(null, com.sun.tools.javac.util.List.nil(),
                memberAccess(classType.toString()),
                com.sun.tools.javac.util.List.nil(), null);
    }

    /**
     * 创建 lambda
     */
    private JCTree.JCVariableDecl createCaller(String varName, JCTree.JCExpression varValue) {
        return JCUtils.createVar(treeMaker,
                names.fromString(varName),
                memberAccess(OriginCaller.class.getName()),
                varValue, Flags.LAMBDA_METHOD);
    }

    /**
     * 递归去扫描
     */
    private JCTree.JCBlock returnNullBody(JCTree.JCBlock block, int reentryCount) {
        ListBuffer<JCTree.JCStatement> result = new ListBuffer<>();

        for (JCTree.JCStatement statement : block.getStatements()) {
            // 如果这个语句是 if 那么继续扫描他的方法块
            if (statement instanceof JCTree.JCIf) {
                JCTree.JCIf jcIf = (JCTree.JCIf) statement;
                jcIf.thenpart = returnNullBody((JCTree.JCBlock) jcIf.thenpart, reentryCount + 1);
                jcIf.elsepart = returnNullBody((JCTree.JCBlock) jcIf.elsepart, reentryCount + 1);
                result.add(jcIf);
                continue;
            }
            //如果是代码块
            if (statement instanceof JCTree.JCBlock) {
                JCTree.JCBlock newJcBlock = returnNullBody((JCTree.JCBlock) statement, reentryCount + 1);
                result.add(newJcBlock);
                continue;
            }

            //如果是 while 语句
            if (statement instanceof JCTree.JCWhileLoop) {
                JCTree.JCWhileLoop jcWhileLoop = (JCTree.JCWhileLoop) statement;
                jcWhileLoop.body = returnNullBody((JCTree.JCBlock) jcWhileLoop.body, reentryCount + 1);
                result.add(jcWhileLoop);
                continue;
            }

            // 如果是 return 语句
            if (statement instanceof JCTree.JCReturn) {
                JCTree.JCReturn jCReturn = (JCTree.JCReturn) statement;
                jCReturn.expr = JCUtils.returnNone(treeMaker);
                result.add(jCReturn);
                continue;
            }

            // 如果是 变量语句，那么复制一遍
            if (statement instanceof JCTree.JCVariableDecl) {
                JCTree.JCVariableDecl jcVariableDecl = (JCTree.JCVariableDecl) statement;
                JCTree.JCVariableDecl newjcVariableDecl = JCUtils.copyVar(treeMaker, jcVariableDecl);
                result.add(newjcVariableDecl);
                continue;
            }
            result.add(statement);
        }
        if (reentryCount == 0) {
            result.add(treeMaker.Return(JCUtils.returnNone(treeMaker)));
        }
        return treeMaker.Block(0, result.toList());
    }

    @Override
    protected Class<? extends Annotation> getNeedToHandleAnnotation() {
        return Proxy.class;
    }
}
